package cn.flying.cloud.gateway.loadbalancer;

import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.stream.Collectors;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.util.CollectionUtils;

import com.alibaba.cloud.nacos.ribbon.ExtendBalancer;
import com.alibaba.cloud.nacos.ribbon.NacosServer;
import com.alibaba.nacos.api.naming.pojo.Instance;
import com.netflix.client.config.IClientConfig;
import com.netflix.loadbalancer.AbstractLoadBalancerRule;
import com.netflix.loadbalancer.DynamicServerListLoadBalancer;
import com.netflix.loadbalancer.Server;

import cn.flying.cloud.base.common.constants.Constants;
import cn.flying.cloud.base.common.utils.lang.ThreadLocalUtil;

/**
 * 基于流量染色的本地负载均衡规则
 *
 * @author: admin
 * @date: 2025年03月16日 11:19
 * @version: 1.0
 */
public class TrafficLoadBalancerRule extends AbstractLoadBalancerRule {
    private final Logger logger = LoggerFactory.getLogger(this.getClass());

    @Override
    public void initWithNiwsConfig(IClientConfig iClientConfig) {
    }

    @Override
    public Server choose(Object key) {
        // 获取服务列表
        String serviceId = ((DynamicServerListLoadBalancer<?>) getLoadBalancer()).getName();
        List<Instance> instanceList =
                getLoadBalancer()
                        .getReachableServers()
                        .stream().map(server -> ((NacosServer) server).getInstance())
                        .collect(Collectors.toList());

        // 当没有实例时快速失败
        if (CollectionUtils.isEmpty(instanceList)) {
            logger.error("[choose] no instance list: {}", serviceId);
            throw new RuntimeException("no instance list");
        }
        // 获取配置
        try {
            if (ThreadLocalUtil.get(Constants.ENV_ID) == null || ThreadLocalUtil.get(Constants.REQUEST_ID) == null) {
                Instance instance = ExtendBalancer.getHostByRandomWeight2(instanceList);
                return new NacosServer(instance);
            }
            String envId = (String) ThreadLocalUtil.get(Constants.ENV_ID);
            if (!"dev".equals(envId)) {
                Instance instance = ExtendBalancer.getHostByRandomWeight2(instanceList);
                return new NacosServer(instance);
            }
            // 根据配置的元数据查询对应的服务，没有则随机选中一个
            String requestId = (String) ThreadLocalUtil.get(Constants.REQUEST_ID);
            Optional<Instance> any = instanceList.stream().filter(instance -> {
                Map<String, String> metadata = instance.getMetadata();
                String metaRequestId = metadata.get("request-id");
                String metaEnvId = metadata.get("env-id");
                if (requestId.equals(metaRequestId) && envId.equals(metaEnvId)) {
                    return true;
                }
                return false;
            }).findAny();
            if (any.isPresent()) {
                return new NacosServer(any.get());
            }
            Instance instance = ExtendBalancer.getHostByRandomWeight2(instanceList);
            return new NacosServer(instance);
        } catch (Exception e) {
            logger.error("[choose] exception: {}", e.getMessage(), e);
            throw new RuntimeException(e);
        }
    }
}